import operator
import os
from typing import Annotated, List

from dotenv import load_dotenv
from langchain_core.messages import BaseMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

class SimpleState(BaseModel):
    """包含用户请求和历史对话的简单状态"""
    query: str
    intent: str=""
    chat_history: Annotated[List[BaseMessage], operator.add]
    log: Annotated[List[str], operator.add]  #累计的日志

# 加载环境变量
load_dotenv()
llm = ChatOpenAI(
    api_key=os.getenv("DASHSCOPE_KEY"),
    base_url=os.getenv("DASHSCOPE_OPENAI_URL"),
    model="qwen-plus-1220"
)


def initial_state(state: SimpleState):
    """初始化全局状态"""
    log_entry = "------初始化状态------"
    print(log_entry)
    return {
        "intent": "",
        "chat_history": [],
        "log": [log_entry]
    }


def triage_node(state: SimpleState):
    """意图识别节点"""
    # 定义提示模板
    prompt_template = ChatPromptTemplate.from_messages([
        ("system", "你是一个意图识别助手。根据用户的输入，判断意图并返回以下之一，注意只能是其中一个英文单词："
                   "- 'chat'：用户进行的是普通聊天。"
                   "- 'knowledge'：用户询问的是Java相关问题。"),
        ("human", "{query}")
    ])
    # 创建链式处理逻辑
    chain = prompt_template | llm | StrOutputParser()
    # 获取用户输入并调用链
    intent = chain.invoke({"query": state.query})

    # 打印日志
    log_entry = f"--- 意图识别结果: {intent} ---"
    print(log_entry)

    # 返回路由决策
    return {
        "intent": intent,
        "log": [log_entry]
    }


def decide_after_triage(state: SimpleState):
    """意图识别之后，决定跳转到哪个节点"""
    if state.intent == "chat":
        return "chat"
    else:
        return "knowledge"